package org.microframework.cloud.gateway.filter.factory;

import java.time.Duration;
import java.time.Instant;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import org.microframework.cloud.gateway.filter.IpKeyResolver;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

/**
 * 基于内存的限流过滤器工厂
 * 使用令牌桶算法实现高效的内存限流功能
 */
@Slf4j
@Component
public class MemoryRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<MemoryRateLimiterGatewayFilterFactory.Config> {

    private final IpKeyResolver ipKeyResolver;
    
    // 存储IP的访问记录
    private final ConcurrentHashMap<String, TokenBucket> buckets = new ConcurrentHashMap<>();
    
    // 桶过期时间（毫秒）
    private static final long BUCKET_EXPIRATION_TIME = 3600000; // 1小时
    
    // 统计信息
    private long totalRequests = 0;
    private long rejectedRequests = 0;
    private long lastStatTime = System.currentTimeMillis();
    
    public MemoryRateLimiterGatewayFilterFactory(IpKeyResolver ipKeyResolver) {
        super(Config.class);
        this.ipKeyResolver = ipKeyResolver;
        
        // 启动桶清理任务
        startBucketCleanupTask();
    }
    
    /**
     * 启动桶清理任务，定期清理过期的令牌桶
     */
    private void startBucketCleanupTask() {
        new Thread(() -> {
            while (true) {
                try {
                    Thread.sleep(60000); // 每分钟执行一次清理
                    cleanExpiredBuckets();
                } catch (InterruptedException e) {
                    log.error("桶清理任务被中断", e);
                    Thread.currentThread().interrupt();
                    break;
                } catch (Exception e) {
                    log.error("桶清理任务异常", e);
                }
            }
        }, "token-bucket-cleanup").start();
    }
    
    /**
     * 清理过期的令牌桶
     */
    private void cleanExpiredBuckets() {
        long now = System.currentTimeMillis();
        int count = 0;
        
        for (String key : buckets.keySet()) {
            TokenBucket bucket = buckets.get(key);
            if (bucket != null && now - bucket.getLastAccessTime() > BUCKET_EXPIRATION_TIME) {
                buckets.remove(key);
                count++;
            }
        }
        
        if (count > 0) {
            log.info("清理了{}个过期的令牌桶，当前剩余{}个", count, buckets.size());
        }
    }
    
    @Override
    public List<String> shortcutFieldOrder() {
        return Arrays.asList("refillRate", "burstCapacity");
    }

    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
            synchronized(this) {
                totalRequests++;
                // 每100000个请求或每小时打印一次统计信息
                if (totalRequests % 100000 == 0 || 
                    System.currentTimeMillis() - lastStatTime > 3600000) {
                    log.info("限流器统计: 总请求={}, 拒绝请求={}, 拒绝率={}%, 活跃桶={}",
                            totalRequests, rejectedRequests, 
                            String.format("%.2f", (double)rejectedRequests/totalRequests*100),
                            buckets.size());
                    lastStatTime = System.currentTimeMillis();
                }
            }
            
            return ipKeyResolver.resolve(exchange)
                .onErrorResume(e -> {
                    log.error("解析IP地址失败", e);
                    return Mono.just("unknown");
                })
                .flatMap(key -> {
                    try {
                        // 获取或创建该IP的令牌桶
                        TokenBucket bucket = buckets.computeIfAbsent(key, 
                                ip -> {
                                    log.debug("为IP={}创建新的令牌桶: 补充率={}, 容量={}", 
                                            ip, config.getRefillRate(), config.getBurstCapacity());
                                    return new TokenBucket(config.getRefillRate(), config.getBurstCapacity());
                                });
                        
                        // 尝试获取令牌
                        if (bucket.tryConsume()) {
                            if (log.isTraceEnabled()) {
                                log.trace("IP={}请求通过: 路径={}, 当前令牌={}", 
                                        key, exchange.getRequest().getPath(), bucket.getTokens());
                            }
                            return chain.filter(exchange)
                                .onErrorResume(e -> {
                                    log.error("限流器处理请求发生错误", e);
                                    return Mono.error(e);
                                });
                        } else {
                            synchronized(this) {
                                rejectedRequests++;
                            }
                            log.warn("请求频率超限: IP={}, 路径={}, 配置={}/{}",
                                    key, exchange.getRequest().getPath(),
                                    config.getRefillRate(), config.getBurstCapacity());
                            
                            return handleRejection(exchange);
                        }
                    } catch (Exception e) {
                        log.error("限流器异常，默认放行: {}", e.getMessage(), e);
                        return chain.filter(exchange);
                    }
                });
        };
    }
    
    /**
     * 处理被拒绝的请求
     */
    private Mono<Void> handleRejection(ServerWebExchange exchange) {
        ServerHttpResponse response = exchange.getResponse();
        response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        response.getHeaders().setContentType(MediaType.APPLICATION_JSON);
        
        String body = "{\"code\": 429, \"message\": \"请求过于频繁，请稍后再试\", \"success\": false}";
        DataBuffer buffer = response.bufferFactory().wrap(body.getBytes());
        return response.writeWith(Mono.just(buffer));
    }
    
    /**
     * 令牌桶实现
     * 基于时间间隔动态补充令牌的高精度实现
     */
    private static class TokenBucket {
        private final double refillRate;  // 每秒填充的令牌数
        private final int capacity;       // 桶容量
        private double tokens;            // 当前令牌数
        private Instant lastRefillTime;   // 上次填充时间
        private long lastAccessTime;      // 上次访问时间
        
        public TokenBucket(double refillRate, int capacity) {
            this.refillRate = refillRate;
            this.capacity = capacity;
            this.tokens = capacity;
            this.lastRefillTime = Instant.now();
            this.lastAccessTime = System.currentTimeMillis();
        }
        
        /**
         * 尝试消费一个令牌
         * @return 是否成功获取到令牌
         */
        public synchronized boolean tryConsume() {
            this.lastAccessTime = System.currentTimeMillis();
            refill();
            
            if (tokens < 1.0) {
                return false;
            }
            
            tokens -= 1.0;
            return true;
        }
        
        /**
         * 获取上次访问时间
         */
        public long getLastAccessTime() {
            return lastAccessTime;
        }
        
        /**
         * 根据时间间隔补充令牌
         */
        private void refill() {
            Instant now = Instant.now();
            double secondsSinceLastRefill = Duration.between(lastRefillTime, now).toMillis() / 1000.0;
            
            // 计算需要补充的令牌数
            double tokensToAdd = secondsSinceLastRefill * refillRate;
            
            if (tokensToAdd > 0) {
                tokens = Math.min(capacity, tokens + tokensToAdd);
                lastRefillTime = now;
            }
        }
        
        /**
         * 获取当前令牌数
         */
        public double getTokens() {
            return tokens;
        }
    }
    
    /**
     * 配置类
     */
    @Data
    public static class Config {
        // 每秒补充的令牌数量
        private double refillRate = 20.0;
        // 令牌桶容量
        private int burstCapacity = 50;
        
        // 支持字符串类型的参数设置
        public void setRefillRate(String refillRate) {
            this.refillRate = Double.parseDouble(refillRate);
        }
        
        public void setBurstCapacity(String burstCapacity) {
            this.burstCapacity = Integer.parseInt(burstCapacity);
        }
    }
} 